'''
@File    : seg2dect.py
@Author  : zym1105
@Date    : 2025-10-10
@Description :
    将分割掩膜 (segmentation mask) 转换为目标检测框 (bbox, normalized xyxy) 的工具函数。
    核心流程：
        1. 连通域提取（使用 OpenCV connectedComponents）
        2. 小框过滤（根据 min_box_size）
        3. 空间聚类与合并（基于矩形间 Chebyshev 距离）
        4. 归一化坐标输出
    适用于点状、多区域的分割结果转检测任务。
'''
from typing import List, Tuple, Optional
from PIL import Image, ImageDraw
import numpy as np
import os
import cv2

def _to_binary_mask(img) -> np.ndarray:
    """Accept PIL.Image or numpy array; return binary mask of shape (H, W) with {0,1}."""
    if isinstance(img, Image.Image):
        arr = np.array(img)
    else:
        arr = np.array(img)
    # 如果是彩色/多通道，取任一通道或转灰度,这里采用简单灰度化，简明处理
    if arr.ndim == 3:
        # 简单灰度化
        arr = (0.2989 * arr[..., 0] + 0.5870 * arr[..., 1] + 0.1140 * arr[..., 2]).astype(np.float32)
    # >0 即前景
    mask = (arr > 0).astype(np.uint8)
    return mask

def _connected_components(mask: np.ndarray, connectivity: int = 8):
    """
    返回 (label_map, num_labels)
    优先用 scipy 或 cv2；都不可用时，用一个简易 BFS 回退实现（足够应付点状/块状目标）。
    """
    # 优先 scipy

    conn = cv2.CV_8U
    connectivity_flag = 8 if connectivity == 8 else 4
    n, labeled = cv2.connectedComponents(mask.astype(np.uint8), connectivity=connectivity_flag)
    return labeled, n - 1  # cv2 的标签从 0..n-1，0 是背景
   

def _boxes_from_label_map(labeled: np.ndarray, num_labels: int,
                          min_long_side: int = 3,
                          min_area: int = 1) -> List[Tuple[int,int,int,int]]:
    """
    从标签图得到候选框，过滤过小框。
    返回像素级 xyxy：[(x1,y1,x2,y2), ...]  (x2,y2 为包含端)
    """
    boxes = []
    H, W = labeled.shape
    for lab in range(1, num_labels+1):
        ys, xs = np.where(labeled == lab)
        if ys.size == 0:
            continue
        y1, y2 = ys.min(), ys.max()
        x1, x2 = xs.min(), xs.max()
        w = x2 - x1 + 1
        h = y2 - y1 + 1
        area = w * h
        if max(w, h) < min_long_side:
            continue
        if area < min_area:
            continue
        boxes.append((x1, y1, x2, y2))
    return boxes

def _chebyshev_rect_distance(a, b) -> int:
    """
    a,b: (x1,y1,x2,y2)  像素坐标，x2,y2 为包含端
    返回两个矩形之间的切比雪夫距离（重叠则为 0）
    """
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    # 计算水平/垂直方向的“间隙”（重叠时为 0）
    dx = max(0, max(ax1 - bx2, bx1 - ax2))
    dy = max(0, max(ay1 - by2, by1 - ay2))
    return max(dx, dy)

def _merge_by_distance(boxes: List[Tuple[int,int,int,int]], thresh_dist: int) -> List[Tuple[int,int,int,int]]:
    """
    以切比雪夫距离 < thresh_dist 为连边，做图的连通分量聚类，并箱。
    """
    n = len(boxes)
    if n == 0:
        return []
    # 建图
    adj = [[] for _ in range(n)]
    for i in range(n):
        for j in range(i+1, n):
            if _chebyshev_rect_distance(boxes[i], boxes[j]) < thresh_dist:
                adj[i].append(j)
                adj[j].append(i)
    # 连通分量
    visited = [False]*n
    merged = []
    for i in range(n):
        if visited[i]: 
            continue
        stack = [i]
        visited[i] = True
        comp = [i]
        while stack:
            u = stack.pop()
            for v in adj[u]:
                if not visited[v]:
                    visited[v] = True
                    stack.append(v)
                    comp.append(v)
        # 组内并箱
        xs1 = [boxes[k][0] for k in comp]
        ys1 = [boxes[k][1] for k in comp]
        xs2 = [boxes[k][2] for k in comp]
        ys2 = [boxes[k][3] for k in comp]
        merged.append((min(xs1), min(ys1), max(xs2), max(ys2)))
    return merged

def _normalize_boxes(boxes: List[Tuple[int,int,int,int]], W: int, H: int) -> List[Tuple[float,float,float,float]]:
    out = []
    for x1, y1, x2, y2 in boxes:
        # 将 x2,y2 视为框外沿，做归一化并裁剪
        nx1 = max(0.0, min(1.0, x1 / float(W)))
        ny1 = max(0.0, min(1.0, y1 / float(H)))
        nx2 = max(0.0, min(1.0, (x2 + 1) / float(W)))  # +1 让闭区间更贴近常见归一化习惯
        ny2 = max(0.0, min(1.0, (y2 + 1) / float(H)))
        # 保证合法顺序
        nx1, nx2 = min(nx1, nx2), max(nx1, nx2)
        ny1, ny2 = min(ny1, ny2), max(ny1, ny2)
        out.append((nx1, ny1, nx2, ny2))

    return out

def seg2dect(
    image_path, 
    threshold_min_box_size: int = 3,
    threshold_min_box_distance: int = 40,
    connectivity: int = 8,
    save_path: Optional[str] = None
) -> List[Tuple[float,float,float,float]]:
    """
    将分割图（mask）转换为合并后的 bbox（归一化 xyxy）
    - image: PIL.Image 或 numpy array；分割前景应为 >0
    - threshold_min_box_size: 小框过滤（以 max(w,h) 度量），像素单位
    - threshold_min_box_distance: 盒间切比雪夫距离阈值（像素）
    - connectivity: 连通域邻域（4 或 8）
    - save_path: 若提供，则把画框可视化结果保存（在原 mask 上画）
    返回：List[ (x1n, y1n, x2n, y2n) ]，均在 [0,1]
    """
    image= Image.open(image_path)
    mask = _to_binary_mask(image)
    H, W = mask.shape
    if H!=512:
        print(image_path)
        print(H,W)
        raise

    # 1) 连通域 → 初始 boxes
    labeled, n = _connected_components(mask, connectivity=connectivity)
    boxes = _boxes_from_label_map(labeled, n, min_long_side=threshold_min_box_size, min_area=1)

    # 2) 距离聚类并箱
    merged_boxes = _merge_by_distance(boxes, threshold_min_box_distance)

    # 3) 归一化
    normed = _normalize_boxes(merged_boxes, W, H)

    # 可视化保存（可选）
    if save_path is not None:
        if isinstance(image, Image.Image):
            vis = image.convert("RGB")
        else:
            # 将 mask 转为灰度可视化背景
            vis = Image.fromarray((mask*255).astype(np.uint8)).convert("RGB")
        draw = ImageDraw.Draw(vis)
        # 原始小框（浅色）
        for x1, y1, x2, y2 in boxes:
            draw.rectangle([x1, y1, x2, y2], outline=(0, 255, 255), width=1)
        # 合并后框（强调色）
        for x1, y1, x2, y2 in merged_boxes:
            draw.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=2)
        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
        vis.save(save_path)

    return normed
